import numpy as np
from scipy.spatial.transform import Rotation
import math
import pinocchio as pin
from os.path import dirname, join


# l_0 l_1 l_5不显含于ik
l_0 = 1
l_1 = 0.135
l_5 = 0.202



l_2=0.5
l_3=0.511


# 坐标D是目标的本体坐标系
DP_C = np.array([0, 0, -l_5])



def iP_2_oP(oR_i, oP_iORG, iP):
    oP = oR_i@iP + oP_iORG

    return oP


def triangle_angles_A(a, b, c):

    # 计算三个角的余弦值
    cos_A = (b**2 + c**2 - a**2) / (2 * b * c)
    
    # 转换为弧度
    A = np.arccos(cos_A)
    
    return A


def eD2oD(eP_D, eD_quat, oT_e = np.array([[0, -1, 0, 0], [1, 0, 0, 0], [0, 0, 1, -l_0-l_1], [0, 0, 0, 1]])):
    oR_e = oT_e[:3, :3]
    oP_eORG = oT_e[:3, 3]
    eR_D = Rotation.from_quat([eD_quat[1], eD_quat[2], eD_quat[3], eD_quat[0]]).as_matrix()
    oR_D = oR_e@eR_D
    oP_D = iP_2_oP(oR_e, oP_eORG, eP_D)
    return oP_D, oR_D



# 坐标o，相当于是1、2电机轴线的交点
def ika1(q_real_float, oP_DORG, oR_D):

    if not hasattr(ika1, 'last_qpos'):  # 初始化前次控制量
        ika1.last_qpos = q_real_float


    # IK_ANALITICAL
    # 还要判断一下各个角度的范围，再据此调整flag
    ############################

    # 求解123
    oP_C = iP_2_oP(oR_D, oP_DORG, DP_C)
    oP_AC = oP_C

    theta_1 = np.arctan2(oP_AC[1], oP_AC[0])

    P_AC_norm = np.linalg.norm(oP_AC)
    Angle_AC_Axis1 = np.arccos(oP_AC[2] / P_AC_norm)
    if P_AC_norm > (l_2 + l_3):
        print("too far")
        return ika1.last_qpos
    elif Angle_AC_Axis1 < 0.1:
        print("123singular")
        return ika1.last_qpos
    else:
        flag = 1
        theta_3 = math.pi - triangle_angles_A(P_AC_norm, l_2, l_3)
        # 求P_{AC}和Axis_1的夹角, arccos的范围是0~pi
        theta_2 = Angle_AC_Axis1-triangle_angles_A(l_3, P_AC_norm, l_2)
    


    # 求解456
    if flag == 1:
        # 计算矩阵^{o}R_{C0}
        oP_B = np.array([l_2*math.sin(theta_2)*math.cos(theta_1),l_2*math.sin(theta_2)*math.sin(theta_1), l_2*math.cos(theta_2)])
        oP_BC = oP_C-oP_B
        oZ_C0 = oP_BC/np.linalg.norm(oP_BC)
        oX_C0 = np.array([-math.sin(theta_1), math.cos(theta_1), 0])
        oY_C0 = np.cross(oZ_C0,oX_C0)
        oR_C0 = np.column_stack((oX_C0, oY_C0, oZ_C0))

        # 计算矩阵^{C0}R_{D}
        C0R_D = oR_C0.T@oR_D
        r = C0R_D
        
        # 计算theta456
        theta_5 = np.arctan2(math.sqrt(r[2][0]**2+r[2][1]**2),r[2][2])
        if theta_5 == 0:
            theta_4 = ika1.last_qpos[3]
            theta_6 = np.arctan2(-r[0][1],r[0][0]) - theta_4
        else:
            theta_4 = np.arctan2(r[1][2]/math.sin(theta_5), r[0][2]/math.sin(theta_5))
            theta_6 = np.arctan2(r[2][1]/math.sin(theta_5), -r[2][0]/math.sin(theta_5))

    
    q_star = np.array([theta_1, theta_2, theta_3, theta_4, theta_5, theta_6])
    ika1.last_qpos = q_star
    return q_star





